'''
Created on Sep 6, 2009
@author: yanghoch
This script provides function to do binary search 
in a sorted table file and output the lines in given range.
'''
import sys
import os
import getopt
import string

class SortedFileSearcher ( object ):
    '''This class can be used to quickly find all lines
     in a increasingly sorted file that has a given int key'''
    def __init__( self, filename ):
        self.fp = open( filename, 'rb' )
        self.fp.seek( 0, 2 ) # set to the end position
        self.length = self.fp.tell() # file length
        self.fp.seek( 0, 0 ) # set back to the first position

    def __get_key__( self, line, colId, delimiter ):
        info = line.split( delimiter )
        if len( info ) > colId:
            return( int( info[colId] ) )
        else:
            print "Unexpected line:" + line
            return( -1 ) # define as dummy key
    def __del__( self ):
        self.fp.close()

    def move2the_begining_of_the_line( self ):
        ''' Move the file pointer to the beginning of
        the line of the given position'''
        fposition = self.fp.tell()
        while fposition > 0:
            self.fp.seek( fposition )
            if self.fp.read( 1 ) == '\n':
                return( self.fp.tell() )
            fposition -= 1
        return 0

    def move2the_begining_of_a_line( self, fposition ):
        ''' Move the file pointer to the beginning of
        the line of the given position'''
        if fposition <= 0:
            fposition = 0
        self.fp.seek( fposition )
        return self.move2the_begining_of_the_line()

    def locatedFp( self, key, colId, delimiter ):
        ''' Pass key,  
            the column num containing the key,  
            and the delimiter to separate columns,
            return the file pointers '''
        low = self.fp.tell()
        high = self.length
        while low < high:
            mid = ( low + high ) // 2
            self.move2the_begining_of_a_line( mid )
            line = self.fp.readline()
            tmp_key = self.__get_key__( line, colId, delimiter )
            # print '--', mid, tmp_key # DEBUG
            if tmp_key < 0:
                break;
            if tmp_key < key:
                low = mid + 1
            else: high = mid
        pivot = low
        self.move2the_begining_of_a_line( pivot )

    def get_lines_w_smaller_key\
        ( self, key, colId, delimiter ):
        ''' Report the lines that have smaller keys
         than the given value after self.fp.
         Assume the file is increasingly sorted.'''
        result = []
        while True:
            line = self.fp.readline()
            if not line: break
            tmp_key = self.__get_key__( line, colId, delimiter )
            if tmp_key < 0: break;
            if tmp_key >= key: break;
            if line[-1:] == '\n': line = line[:-1] # remove '\n'
            result.append( line )
        return( result )

    def get_lines_w_greater_key\
        ( self, key, colId, delimiter ):
        ''' Report the lines that have larger keys
         than the given value after self.fp. 
         Assume the file is increasingly sorted.'''
        result = []
        while True:
            line = self.fp.readline()
            if not line: break
            tmp_key = self.__get_key__( line, colId, delimiter )
            if tmp_key < 0: break;
            if tmp_key > key:
                if line[-1:] == '\n': line = line[:-1]
            result.append( line )
        return( result )

def extract_record( iFileN, oFileN, posColId, range, shift ):
    '''Given the table with all sorted coverage info,
    get all the fragment coverage info for a given range.
    @param args[0]: input filename
    @param args[1]: output filename 
    @param args[2]: lower bound of the extracted regions (left end).
    @param args[3]: upper bound of the extracted regions (left end).
    @param args[4]: The shift for printing the index.'''
    delimiter = '\t'
    if( not os.path.exists( iFileN ) ):
        print iFileN, " doesn't exist."
        return False
    s = SortedFileSearcher( iFileN );
    oFile = open( oFileN, "w" );
    startPos = min( range )
    endPos = max( range )
    s.locatedFp( startPos, posColId, delimiter )
    records = s.get_lines_w_smaller_key( endPos, posColId , delimiter )
    # sys.stderr.write( "There are " + str( len( records ) ) + " regions" )

    for line in records:
        info = line.split( delimiter )
        try:
            info[posColId] = str( int( info[posColId] ) - int( shift ) )
            info[posColId + 1] = str( int( info[posColId + 1] ) - int( shift ) )
            oFile.write( string.join( info, '\t' ) + "\n" )
        except ValueError:
            sys.stderr = open( "ValueError.log", "w" )
            sys.stderr.write( "\nValueError in " + iFileN )
            sys.stderr.write( "\n" + line )
            sys.stderr.write( "\nposColId " + str( posColId ) )
            sys.stderr.write( "\ninfo[posColId] " + str( info[posColId] ) )
            sys.stderr.write( "\ninfo[posColId + 1] " + str( info[posColId + 1] ) )
            sys.stderr.write( "\nshift " + str( shift ) + "\n" )
            sys.stderr = sys.__stderr__
    oFile.close()

def usage():
    sys.stderr.write( "Working dir is " + os.getcwd() + '\n' )
    print __doc__
    print "Four arguments are required."
    print "( 1 ) Input file"
    print "( 2 ) Output file"
    print "( 3 ) range cols ID (Range are assumed in a consecutive columns"
    print "( 4 ) start range"
    print "( 5 ) end range"
    print "( 6 ) the optional argument is the shift of index to output."

def main():
    try:
        opts, args = getopt.getopt( sys.argv[1:], 'h', ["help"] )
    except getopt.error, msg:
        print msg
        sys.exit( 2 );
    #process options and argument 
    for o in opts:
        if o in ( "-h", "--help" ):
            print __doc__
            sys.exit();
    if ( len( args ) == 5 or len( args ) == 6 ):
        ifileN = args[0]
        ofileN = args[1]
        colId = int( args[2] )
        range = [ int( args[3] ), int( args[4] ) ]
        if len( args ) == 5:
            shift = 0
        else:
            shift = int( args[5] )
        if os.path.exists( ifileN ):
            extract_record( ifileN, ofileN, colId, range, shift )
        else:
            print ifileN, " doesn't exist."
    else:
        usage()
if __name__ == '__main__':
    main()
